from contextlib import contextmanager

from sqlalchemy import URL
from sqlmodel import Session, create_engine
from config.env import MysqlConfig
# from entity.model.test_model import *
# from entity.model.crawler_record_model import *
# from entity.model.hefei.crawler_record_model import *
# from entity.model.community_model import *
from util.mylog import my_logger

# mysql
mysql_engine = create_engine(
    url=URL.create(
        drivername="mysql+pymysql",
        username=MysqlConfig.db_username,
        password=MysqlConfig.db_password,
        host=MysqlConfig.db_host,
        port=MysqlConfig.db_port,
        database=MysqlConfig.db_database,
        query={
            "autocommit": "true",
        },
    ),
    echo=MysqlConfig.db_echo,  # 查看SQL语句日志
    pool_size=MysqlConfig.db_pool_size,  # 设置连接池的大小
    max_overflow=MysqlConfig.db_max_overflow,  # 设置最大溢出连接数
    pool_recycle=MysqlConfig.db_pool_recycle,  # 设置连接回收时间（秒）
    pool_timeout=MysqlConfig.db_pool_timeout,  # 设置连接超时时间（秒）
)
@contextmanager
def get_session():
    with Session(bind=mysql_engine) as session:
        try:
            yield session
        except Exception as e:
            my_logger.error(f"mysql session error: {e}")
            session.rollback()
            # 直接raise而不是raise e，以保留原始的异常追踪信息
            raise
        finally:
            session.close()

# SQLModel.metadata.create_all(mysql_engine)